# #!/bin/bash

gpu=0
dataset_dir='./data'

# ## This super-network will be generated after executing 1_script_train_supernet.sh
# #search_mode=random ## WS + random search (in Table 5)
# search_mode=evolution ## AutoSNN
dataset_name=CIFAR10
T=8 #20
suffix_prefix='_SNN_Adam_600ep_2022'
batch_size=96 #64
for channels in 16 32 128
do
# searchString="DVS"
# for dataset_name in CIFAR10DVS DVS128Gesture Tiny-ImageNet-200
# do

# isDVS=$(echo $dataset_name | grep "${searchString}")

# if [[ "$isDVS" != "" ]] 
# then
#     echo "$dataset_name is DVS dataset"
#     T=20
#     channels=16
#     suffix_prefix='_Adam_2022_T_'$T'_init_tau_2.0_vth_1.0_neuron_PLIF_split_by_number_normalization_None'
#     batch_size=64
# else
#     echo "$dataset_name is not DVS dataset"
#     T=8 #20
#     channels=64 #16
#     suffix_prefix='_SNN_Adam_600ep_2022'
#     batch_size=96 #64
# fi

suffix=$suffix_prefix'/batch_size_'$batch_size'/checkpoint.pth.tar'
search_space=AutoSNN_$channels
prefix='macro_search_result2/uniform_sampling/'$search_space'_'
trained_supernet=$prefix$dataset_name$suffix

python train_supernet/train.py \
    --gpu $gpu \
    --T $T --init_tau 2.0 --v_threshold 1.0 --neuron PLIF \
    --epochs 600 \
    --dataset_dir $dataset_dir \
    --dataset_name $dataset_name \
    --save uniform_sampling \
    --search_space $search_space \
    --seed 2022 \
    --batch_size $batch_size
    
echo "Training supernet on $dataset_name is completed \n"

for search_mode in random evolution
do
python search_arch/search.py \
    --gpu $gpu \
    --T $T --init_tau 2.0 --v_threshold 1.0 --neuron PLIF \
    --dataset_dir $dataset_dir \
    --dataset_name $dataset_name \
    --supernet $trained_supernet \
    --seed 2022 \
    --search_space $search_space \
    --search_algo $search_mode \
    --fitness ACC_pow_spikes \
    --fitness_lambda -0.08 \
    --batch_size $batch_size

echo "Searching AutoSNN on $dataset_name with $search_mode is completed \n"

done
done